# cython: profile=True
import cython
from numpy cimport ndarray
from v3dsolver.fluid_elements.fluid_elems cimport FluidElem,Saffman,PseudoParticle
from v3dsolver.tree_manager.tree_manager cimport Cell,MakinoCell,TreeManager
from v3dsolver.vector cimport Vec,Vec_new
from mpi4py cimport MPI

cdef class VelCalculator:
    cdef public TreeManager tree
    cdef public MPI.Comm comm
    
    @cython.locals(effect=FluidElem,cause=FluidElem,v=Vec)
    cpdef direct_vel_computation(self,elems)

    @cython.locals(elem=FluidElem)
    cpdef serial_vel_computation(self)

    @cython.locals(elem=FluidElem,root=MakinoCell,v=Vec)
    cdef traverse_own_tree(self)

    @cython.locals(celldata=list)
    cdef gather_celldata(self,int level)

    @cython.locals(nearby=list,faroff=list,maxlevel=int,child=Cell)
    cdef mark_nearby_faroff(self,Cell cell,Vec cellcenter,int level)

    #@cython.locals(myrank=int,send_objects=list,i=int,destrank=int,j=int,celldata=list,
    #               cellcenter=Vec,cellid=int,nearby=list,faroff=list,cell=Cell,nearby_ids=list,
    #               pseudo_pos_vors=list,pseudo=PseudoParticle,cells_data=list,sendobj=list,
    #               all_nearby_faroff_data=list)
    #cdef gather_nearby_faroff_data(self,list all_celldata,int level)

    @cython.locals(i=int,celldata=list,cellid=int,pseudo_pos_vors=list,
                   cell=Cell,pseudop=Saffman,pseudo_pos_vor=list,elem=FluidElem,
                   v=Vec)
    cdef compute_faroff_vel(self,list data)

    @cython.locals(sendback_particles=list,i=int,celldata=list,cellid=int,nearby_ids=list,cell=Cell,elem=FluidElem,
                   elem_positions=list)
    cdef send_particles(self,list data,int server_rank)

    @cython.locals(j=int,celldata=list,cellid=int,nearby_ids=list,elem_positions=list,id=int,
                   nearby_cells=list,i=int,vels=list,k=int,elem_pos=Vec,cause_cell=MakinoCell,
                   dummy_elem=FluidElem,send_vels_obj=list,v=Vec)
    cdef vel_recvd_particles(self,list recvd_data,int client_rank)

    @cython.locals(j=int,celldata=list,cellid=int,velocities=list,cell=Cell,i=int,elem=FluidElem)
    cdef assign_nearby_vels(self,list recvd_vel_data)

    cdef share_data(self,list data)

    @cython.locals(level=int,all_celldata=list,all_nearby_faroff_data=list,particle_data=list,
                   nearfar_data=list,server=int,client=int,src_data=list,recvd_particle_data=list,
                   vel_data=list,)
    cpdef parallel_vel_computation(self)
